rm(list = ls())

set.seed(1)

library(minqa)
library(VGAM)
library(gtools)

setwd("~/Dropbox/complieropt/Polmeth Archive")
source("icsw.R")

###

dgp <- function(Z,X1,X2,X3,a,b,c,d,e,f,g,h) {
	N <- length(Z)
	AC <- (2 + a + b*X1 + c*X2 + d*X3 + rnorm(N,0,3) > 0)
	A <- (-2 + e + f*X1 + g*X2 + h*X3 + rnorm(N,0,3) > 0)*AC
	C <- AC - A
	
	D <- C*Z + A
	Y0 <- 5*AC + rnorm(N)
	tau<- 5*(1+X1+X2+X3)
	Y <- Y0 + tau*(D)
	
	return(list(D,Y,tau,C))
	}

perm.test <- function(D,Z,Xmat,cscoreout,numperm=50,genoud=FALSE) {
	test.dist <- rep(NA,numperm)
	
	for (i in 1:numperm) {
		Xmats <- Xmat[sample(1:length(D)),]
		cout <- complier.score(D,Z,Xmats,genoud=genoud)
		test.dist[i] <- sum((D - cout$C.score*Z - cout$A.score)^2)
	}
	
	test.stat <- sum((D - cscoreout$C.score*Z - cscoreout$A.score)^2)
	p.value <- mean(test.dist <= test.stat)
	output <- list(test.dist,test.stat,p.value)
	names(output) <- c("test.dist","test.stat","p.value")
	return(output)
}

# ATE is always 1

iteration <- function(N) {

alpha <- 0.275

Z <- sample(c(rep(1,N/2),rep(0,N/2)))

X1 <- runif(N,-1,1)
X2 <- runif(N,-1,1)
X3 <- runif(N,-1,1)

a <- runif(1,-2,2)
b <- runif(1,-2,2)
c <- runif(1,-2,2)
d <- runif(1,-2,2)
e <- runif(1,-2,2)
f <- runif(1,-2,2)
g <- runif(1,-2,2)
h <- runif(1,-2,2)

output <- rep(NA,8)

# DGP1: Correct Model

dgpout <- dgp(Z,X1,X2,X3,a,b,c,d,e,f,g,h)

D <- dgpout[[1]]
Y <- dgpout[[2]]

Xmat <- W <- cbind(X1,X2,X3)

cscoreout <- complier.score(genoud=FALSE,D,Z,Xmat)
cscore <- cscoreout$C.score
qcscore <- quantile(cscore,1/N^alpha)
cscore[cscore < qcscore] <- qcscore

output[1] <- mean(dgpout[[3]])
output[2] <- lm(Y~Z+Xmat)$coefficients["Z"]
output[3] <- lm(Y~D+Xmat)$coefficients["D"]
output[4] <- tsls.wfit(cbind(1,D,Xmat),Y,cbind(1,Z,Xmat),weights=rep(1,N))$coefficients["D"]
output[5] <- tsls.wfit(cbind(1,D,Xmat),Y,cbind(1,Z,Xmat),weights=1/cscore)$coefficients["D"]
output[6] <- mean(dgpout[[4]])
output[7] <- perm.test(D,Z,Xmat,cscoreout,numperm=1)$p.value
output[8] <- N

return(output)
}

numiter <- 2500

#N: 500, 1000, 2500, 5000, 10000, 25000

Nlist <- c(500, 1000, 2500, 5000, 10000, 25000)

system.time(simoutput <- sapply(sort(rep(Nlist,numiter)),iteration))

save.image(paste("Right.rdata",sep=""))

simoutput[8,] == 500 -> isN
t500 <- c(colMeans(apply(simoutput[2:5,isN],1,function(x) (x - simoutput[1,isN])^2))^.5,mean(simoutput[7,isN]))
simoutput[8,] == 1000 -> isN
t1000 <- c(colMeans(apply(simoutput[2:5,isN],1,function(x) (x - simoutput[1,isN])^2))^.5,mean(simoutput[7,isN]))
simoutput[8,] == 2500 -> isN
t2500 <- c(colMeans(apply(simoutput[2:5,isN],1,function(x) (x - simoutput[1,isN])^2))^.5,mean(simoutput[7,isN]))
simoutput[8,] == 5000 -> isN
t5000 <- c(colMeans(apply(simoutput[2:5,isN],1,function(x) (x - simoutput[1,isN])^2))^.5,mean(simoutput[7,isN]))
simoutput[8,] == 10000 -> isN
t10000 <- c(colMeans(apply(simoutput[2:5,isN],1,function(x) (x - simoutput[1,isN])^2))^.5,mean(simoutput[7,isN]))
simoutput[8,] == 25000 -> isN
t25000 <- c(colMeans(apply(simoutput[2:5,isN],1,function(x) (x - simoutput[1,isN])^2))^.5,mean(simoutput[7,isN]))

round(rbind(t500,t1000,t2500,t5000,t10000,t25000),3)



